import numpy as np
import matplotlib.pyplot as plt
import random
import time
import itertools
from collections import defaultdict

# ================== 参数配置区 ==================
# 用户可自由修改以下参数
CITY_NUM = 20  # 城市数量
POP_SIZE = 100  # 种群规模（个体数量）
MAX_GENERATIONS = 500  # 最大迭代次数
MUTATION_RATE = 0.2  # 变异概率
CROSSOVER_RATE = 0.8  # 交叉概率
ELITE_SIZE = 10  # 精英个体数量（直接保留到下一代）


# ================== 城市生成与距离计算 ==================
def generate_cities(num_cities):
    """生成随机城市坐标"""
    np.random.seed(100)  # 固定随机种子，确保结果可复现
    return np.random.rand(num_cities, 2) * 100


def calculate_distance(city1, city2):
    """计算两个城市之间的欧几里得距离"""
    return np.sqrt((city1[0] - city2[0]) ** 2 + (city1[1] - city2[1]) ** 2)


def calculate_total_distance(path, cities):
    """计算路径总距离"""
    total_distance = 0
    num_cities = len(path)

    for i in range(num_cities):
        from_city = path[i]
        to_city = path[(i + 1) % num_cities]  # 循环回到起点
        total_distance += calculate_distance(cities[from_city], cities[to_city])

    return total_distance


# ================== 遗传算法操作 ==================
def initialize_population(pop_size, num_cities):
    """初始化种群：生成随机路径"""
    population = []
    for _ in range(pop_size):
        # 生成0到num_cities-1的随机排列
        path = list(range(num_cities))
        random.shuffle(path)
        population.append(path)
    return population


def rank_population(population, cities):
    """评估种群并排序（按路径长度升序）"""
    fitness_results = []
    for i, path in enumerate(population):
        distance = calculate_total_distance(path, cities)
        fitness_results.append((i, distance))

    # 按路径长度排序（距离越短越好）
    return sorted(fitness_results, key=lambda x: x[1])


# ================== 选择操作 ==================
def roulette_wheel_selection(population, ranked_population, elite_size):
    """轮盘赌选择"""
    selection_results = []

    # 1. 首先保留精英个体
    for i in range(elite_size):
        selection_results.append(population[ranked_population[i][0]])

    # 2. 计算适应度值（距离的倒数）
    fitness_values = [1 / ranked_population[i][1] for i in range(len(ranked_population))]
    total_fitness = sum(fitness_values)

    # 3. 计算选择概率
    selection_probs = [f / total_fitness for f in fitness_values]

    # 4. 使用轮盘赌选择剩余个体
    for _ in range(len(population) - elite_size):
        pick = random.random()
        current = 0
        for i in range(len(ranked_population)):
            current += selection_probs[i]
            if current > pick:
                selection_results.append(population[ranked_population[i][0]])
                break

    return selection_results


def tournament_selection(population, ranked_population, elite_size, tournament_size=5):
    """锦标赛选择"""
    selection_results = []

    # 1. 首先保留精英个体
    for i in range(elite_size):
        selection_results.append(population[ranked_population[i][0]])

    # 2. 锦标赛选择剩余个体
    for _ in range(len(population) - elite_size):
        tournament = random.sample(ranked_population, tournament_size)
        winner = min(tournament, key=lambda x: x[1])  # 选择距离最短的个体
        selection_results.append(population[winner[0]])

    return selection_results


# ================== 交叉操作 ==================
def ordered_crossover(parent1, parent2):
    """有序交叉(OX)"""
    size = len(parent1)

    # 选择两个随机交叉点
    start, end = sorted(random.sample(range(size), 2))

    # 初始化子代
    child = [None] * size

    # 从parent1复制片段
    child[start:end] = parent1[start:end]

    # 从parent2填充剩余部分
    current_pos = end % size
    for gene in parent2:
        if gene not in child:
            child[current_pos] = gene
            current_pos = (current_pos + 1) % size

    return child


def partially_mapped_crossover(parent1, parent2):
    """修正后的部分映射交叉(PMX)"""
    size = len(parent1)

    # 确保交叉点不相同
    while True:
        start, end = sorted(random.sample(range(size), 2))
        if end - start > 1:  # 确保有足够的交叉长度
            break

    # 初始化子代
    child1 = [None] * size
    child2 = [None] * size

    # 复制交叉段
    child1[start:end] = parent1[start:end]
    child2[start:end] = parent2[start:end]

    # 创建映射关系
    mapping1 = {}
    mapping2 = {}

    for i in range(start, end):
        if parent2[i] not in parent1[start:end]:
            mapping1[parent2[i]] = parent1[i]
        if parent1[i] not in parent2[start:end]:
            mapping2[parent1[i]] = parent2[i]

    # 填充子代1
    for i in list(range(0, start)) + list(range(end, size)):
        gene = parent2[i]
        while gene in mapping1:
            gene = mapping1[gene]
        child1[i] = gene

    # 填充子代2
    for i in list(range(0, start)) + list(range(end, size)):
        gene = parent1[i]
        while gene in mapping2:
            gene = mapping2[gene]
        child2[i] = gene

    return child1, child2


def validate_path(path, num_cities):
    """验证路径是否合法"""
    if len(path) != num_cities:
        return False
    if len(set(path)) != num_cities:
        return False
    if min(path) < 0 or max(path) >= num_cities:
        return False
    return True


# ================== 变异操作 ==================
def swap_mutation(path):
    """交换变异：随机交换两个城市的位置"""
    if random.random() < MUTATION_RATE:
        idx1, idx2 = random.sample(range(len(path)), 2)
        path[idx1], path[idx2] = path[idx2], path[idx1]
    return path


def inversion_mutation(path):
    """逆序变异：随机选择一段路径并反转"""
    if random.random() < MUTATION_RATE:
        start, end = sorted(random.sample(range(len(path)), 2))
        path[start:end] = list(reversed(path[start:end]))
    return path


def scramble_mutation(path):
    """打乱变异：随机选择一段路径并打乱"""
    if random.random() < MUTATION_RATE:
        start, end = sorted(random.sample(range(len(path)), 2))
        segment = path[start:end]
        random.shuffle(segment)
        path[start:end] = segment
    return path


# ================== 遗传算法主函数 ==================
def genetic_algorithm(cities, selection_method, crossover_method, mutation_method):
    """修正后的遗传算法主函数"""
    # 记录每代最佳距离
    progress = []
    convergence_gen = 0
    best_distance_history = []
    num_cities = len(cities)
    start_time = time.time()

    # 1. 初始化种群
    population = initialize_population(POP_SIZE, num_cities)

    # 2. 评估初始种群
    ranked_population = rank_population(population, cities)
    best_distance = ranked_population[0][1]
    progress.append(best_distance)
    best_distance_history.append(best_distance)

    # 3. 进化过程
    for generation in range(MAX_GENERATIONS):
        # 3.1 选择
        selected = selection_method(population, ranked_population, ELITE_SIZE)

        # 3.2 交叉
        children = []
        # 首先保留精英
        children.extend(selected[:ELITE_SIZE])

        # 交叉生成后代
        for i in range(ELITE_SIZE, len(selected), 2):
            if i + 1 >= len(selected):
                break  # 防止索引越界

            parent1 = selected[i]
            parent2 = selected[i + 1]

            if random.random() < CROSSOVER_RATE:
                # 根据交叉方法生成子代
                if crossover_method == partially_mapped_crossover:
                    child1, child2 = crossover_method(parent1, parent2)
                    # 验证子代
                    if validate_path(child1, num_cities) and validate_path(child2, num_cities):
                        children.append(child1)
                        children.append(child2)
                    else:
                        # 无效子代，使用父代代替
                        children.append(parent1)
                        children.append(parent2)
                else:
                    child = crossover_method(parent1, parent2)
                    if validate_path(child, num_cities):
                        children.append(child)
                    else:
                        children.append(parent1)

                    child2 = crossover_method(parent2, parent1)
                    if validate_path(child2, num_cities):
                        children.append(child2)
                    else:
                        children.append(parent2)
            else:
                # 不交叉，直接复制父代
                children.append(parent1)
                children.append(parent2)

        # 3.3 变异
        for i in range(len(children)):
            original = children[i][:]  # 保存副本
            children[i] = mutation_method(children[i])
            if not validate_path(children[i], num_cities):
                children[i] = original  # 变异无效，恢复原路径

        # 3.4 更新种群
        population = children

        # 3.5 评估新种群
        ranked_population = rank_population(population, cities)
        current_best = ranked_population[0][1]
        progress.append(current_best)

        # 更新历史最佳
        if current_best < best_distance:
            best_distance = current_best
            convergence_gen = generation
        best_distance_history.append(best_distance)

        # 3.6 打印进度
        if generation % 50 == 0:
            print(f"Generation {generation}: Best Distance = {current_best:.2f}")

    end_time = time.time()

    # 4. 返回最终结果
    best_index = ranked_population[0][0]
    best_path = population[best_index]
    best_distance = ranked_population[0][1]

    print(f"\nFinal Result: Best Distance = {best_distance:.2f}")
    print(f"Converged at generation: {convergence_gen}")
    print(f"Execution time: {end_time - start_time:.2f} seconds")

    return best_path, best_distance, progress, convergence_gen, best_distance_history


# ================== 可视化函数 ==================
def plot_cities(cities):
    """绘制城市分布图"""
    plt.figure(figsize=(10, 6))
    x = [city[0] for city in cities]
    y = [city[1] for city in cities]
    plt.scatter(x, y, c='red', marker='o')

    for i, city in enumerate(cities):
        plt.annotate(str(i), (city[0], city[1]), fontsize=12)

    plt.title(f'City Distribution ({len(cities)} Cities)')
    plt.xlabel('X Coordinate')
    plt.ylabel('Y Coordinate')
    plt.grid(True)
    plt.show()


def plot_path(path, cities, title):
    """绘制路径图"""
    plt.figure(figsize=(10, 6))
    x = [city[0] for city in cities]
    y = [city[1] for city in cities]
    plt.scatter(x, y, c='red', marker='o')

    # 绘制路径
    for i in range(len(path)):
        start_city = cities[path[i]]
        end_city = cities[path[(i + 1) % len(path)]]
        plt.plot([start_city[0], end_city[0]], [start_city[1], end_city[1]], 'b-')

    # 标记起点
    start_city = cities[path[0]]
    plt.annotate('Start', (start_city[0], start_city[1]), fontsize=12, color='green')

    plt.title(f'{title} (Distance: {calculate_total_distance(path, cities):.2f})')
    plt.xlabel('X Coordinate')
    plt.ylabel('Y Coordinate')
    plt.grid(True)
    plt.show()


def plot_progress(progress, title):
    """绘制进化过程"""
    plt.figure(figsize=(10, 6))
    plt.plot(progress)
    plt.title(title)
    plt.xlabel('Generation')
    plt.ylabel('Best Distance')
    plt.grid(True)
    plt.show()


def plot_comparison(results):
    """绘制不同算法组件组合的结果比较"""
    plt.figure(figsize=(14, 8))

    # 创建颜色和线型循环
    colors = plt.cm.tab10(np.linspace(0, 1, len(results)))
    line_styles = ['-', '--', '-.', ':']

    for i, (config, data) in enumerate(results.items()):
        # 提取配置名称
        config_name = config

        # 绘制收敛曲线
        plt.plot(data['progress'],
                 label=f"{config_name} (Dist: {data['best_distance']:.2f}, Conv: {data['convergence_gen']})",
                 color=colors[i % len(colors)],
                 linestyle=line_styles[(i // len(colors)) % len(line_styles)],
                 linewidth=2)

    plt.title('Comparison of Different Genetic Algorithm Components')
    plt.xlabel('Generation')
    plt.ylabel('Best Distance')
    plt.legend(fontsize=10)
    plt.grid(True)
    plt.tight_layout()
    plt.show()


def plot_component_comparison(component_results, component_type):
    """绘制特定组件的比较结果"""
    plt.figure(figsize=(14, 8))

    # 创建颜色和线型循环
    colors = plt.cm.tab10(np.linspace(0, 1, len(component_results)))

    for i, (method, result) in enumerate(component_results.items()):
        # 绘制收敛曲线
        plt.plot(result['progress'],
                 label=f"{method} (Dist: {result['best_distance']:.2f}, Conv: {result['convergence_gen']})",
                 color=colors[i],
                 linewidth=2)

    plt.title(f'Comparison of Different {component_type} Methods')
    plt.xlabel('Generation')
    plt.ylabel('Best Distance')
    plt.legend(fontsize=10)
    plt.grid(True)
    plt.tight_layout()
    plt.show()


def print_results_table(results):
    """打印结果表格"""
    print("\n" + "=" * 80)
    print("Genetic Algorithm Components Comparison Results")
    print("=" * 80)
    print(f"{'Configuration':<30} | {'Best Distance':<15} | {'Convergence Gen':<15} | {'Description'}")
    print("-" * 80)

    # 默认配置
    default = results['Default']
    print(
        f"{'Default':<30} | {default['best_distance']:<15.2f} | {default['convergence_gen']:<15} | Tournament + OX + Swap")

    # 选择算子比较
    roulette = results['Roulette Selection']
    print(
        f"{'Roulette Selection':<30} | {roulette['best_distance']:<15.2f} | {roulette['convergence_gen']:<15} | Roulette wheel selection")

    # 交叉算子比较
    pmx = results['PMX Crossover']
    print(
        f"{'PMX Crossover':<30} | {pmx['best_distance']:<15.2f} | {pmx['convergence_gen']:<15} | Partially Mapped Crossover")

    # 变异算子比较
    inversion = results['Inversion Mutation']
    scramble = results['Scramble Mutation']
    print(
        f"{'Inversion Mutation':<30} | {inversion['best_distance']:<15.2f} | {inversion['convergence_gen']:<15} | Inversion mutation")
    print(
        f"{'Scramble Mutation':<30} | {scramble['best_distance']:<15.2f} | {scramble['convergence_gen']:<15} | Scramble mutation")

    print("=" * 80)


# ================== 主程序 ==================
if __name__ == "__main__":
    # 生成城市
    cities = generate_cities(CITY_NUM)
    plot_cities(cities)

    # 定义不同算法组件组合
    # 定义实验配置字典，用于存储不同的遗传算法组件组合
    experiments = {
        # 默认配置：锦标赛选择 + 有序交叉(OX) + 交换变异
        "Default": {
            "selection": tournament_selection,  # 使用锦标赛选择算子
            "crossover": ordered_crossover,  # 使用有序交叉(OX)算子
            "mutation": swap_mutation,  # 使用交换变异算子
            "description": "Tournament + OX + Swap"  # 配置描述
        },

        # 轮盘赌选择实验：比较不同选择算子的影响
        "Roulette Selection": {
            "selection": roulette_wheel_selection,  # 使用轮盘赌选择算子（替换默认的锦标赛选择）
            "crossover": ordered_crossover,  # 保持有序交叉不变
            "mutation": swap_mutation,  # 保持交换变异不变
            "description": "Roulette wheel selection"  # 配置描述
        },

        # PMX交叉实验：比较不同交叉算子的影响
        "PMX Crossover": {
            "selection": tournament_selection,  # 保持锦标赛选择不变
            "crossover": partially_mapped_crossover,  # 使用部分映射交叉(PMX)算子（替换默认的有序交叉）
            "mutation": swap_mutation,  # 保持交换变异不变
            "description": "Partially Mapped Crossover"  # 配置描述
        },

        # 逆序变异实验：比较不同变异算子的影响
        "Inversion Mutation": {
            "selection": tournament_selection,  # 保持锦标赛选择不变
            "crossover": ordered_crossover,  # 保持有序交叉不变
            "mutation": inversion_mutation,  # 使用逆序变异算子（替换默认的交换变异）
            "description": "Inversion mutation"  # 配置描述
        },

        # 打乱变异实验：比较不同变异算子的影响
        "Scramble Mutation": {
            "selection": tournament_selection,  # 保持锦标赛选择不变
            "crossover": ordered_crossover,  # 保持有序交叉不变
            "mutation": scramble_mutation,  # 使用打乱变异算子（替换默认的交换变异）
            "description": "Scramble mutation"  # 配置描述
        }
    }

    # 存储所有结果
    results = {}
    selection_results = {}
    crossover_results = {}
    mutation_results = {}

    print(f"=== 遗传算法参数 ===")
    print(f"城市数量: {CITY_NUM}")
    print(f"种群大小: {POP_SIZE}")
    print(f"最大代数: {MAX_GENERATIONS}")
    print(f"变异概率: {MUTATION_RATE}")
    print(f"交叉概率: {CROSSOVER_RATE}")
    print(f"精英数量: {ELITE_SIZE}")

    # 运行所有实验
    for name, config in experiments.items():
        print(f"\n{'=' * 50}")
        print(f"Running experiment: {name}")
        print(f"Configuration: {config['description']}")
        print(f"{'=' * 50}")

        # 运行遗传算法
        best_path, best_distance, progress, convergence_gen, best_history = genetic_algorithm(
            cities,
            selection_method=config['selection'],
            crossover_method=config['crossover'],
            mutation_method=config['mutation']
        )

        # 保存结果
        results[name] = {
            "best_path": best_path,
            "best_distance": best_distance,
            "progress": progress,
            "convergence_gen": convergence_gen,
            "best_history": best_history,
            "description": config['description']
        }

        # 分类保存结果用于单独比较
        if "Roulette" in name:
            selection_results[name] = results[name]
        elif "PMX" in name:
            crossover_results[name] = results[name]
        elif "Inversion" in name or "Scramble" in name:
            mutation_results[name] = results[name]

    # 打印结果表格
    print_results_table(results)

    # 绘制所有结果的比较图
    plot_comparison(results)

    # 绘制选择算子的比较
    if selection_results:
        plot_component_comparison(selection_results, "Selection Methods")

    # 绘制交叉算子的比较
    if crossover_results:
        plot_component_comparison(crossover_results, "Crossover Methods")

    # 绘制变异算子的比较
    if mutation_results:
        plot_component_comparison(mutation_results, "Mutation Methods")

    # 绘制最优路径
    for name, result in results.items():
        plot_path(result['best_path'], cities, f"{name} - {result['description']}")